package com.tianji.aigc.memory;

import cn.hutool.core.bean.BeanUtil;
import cn.hutool.core.collection.CollUtil;
import cn.hutool.json.JSONUtil;
import com.tianji.aigc.entity.RedisMessage;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.*;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.stereotype.Component;

import java.util.List;

/**
 * TODO
 *
 * @Author mr.xing
 * @Date 2025/9/2 21:27
 */
@Component
public class RedisChatMemory implements ChatMemory {

    @Autowired
    private RedisTemplate redisTemplate;

    /**
     *  保存会话对应的消息到redis中
     * @param conversationId
     * @param messages
     */
    @Override
    public void add(String conversationId, List<Message> messages) {
        if (CollUtil.isEmpty(messages)){
            return;
        }
        // 获取rediskey,组成结构 : CHAT:用户id:sessionID
        String key = getRediskey (conversationId);
        // 将消息转为字符串类型
        List<String> stringList = messages.stream().map(message -> {
                    //拷贝Message对象为自定义的RedisMessage对象
                    RedisMessage redisMessage = BeanUtil.toBean(message, RedisMessage.class);
                    //设置消息内容字符串，否则为空
                    redisMessage.setTextContent(message.getText());
                    //设置Tool相关属性的值
                    if(message instanceof AssistantMessage assistantMessage) {
                        redisMessage.setToolCalls(assistantMessage.getToolCalls());
                    } else if(message instanceof ToolResponseMessage toolResponseMessage) {
                        redisMessage.setToolResponses(toolResponseMessage.getResponses());
                    }
                 return JSONUtil.toJsonStr(redisMessage);
        }).toList();
        // 将消息存入Redis的List中
        redisTemplate.opsForList().leftPushAll(key, stringList);
    }

    /**
     * 获取对话ID相关的消息列表
     * @param conversationId
     * @param lastN
     * @return
     */
    @Override
    public List<Message> get(String conversationId, int lastN) {
        // 拼接Rediskey
        String key = getRediskey (conversationId);
        // 获取消息字符串列表
        List<String> range = redisTemplate.opsForList().range(key, 0, lastN);
        // 将消息字符串转为消息对象
        List<Message> list = range.stream().map(m -> {
                    //将消息JSON字符串转为RedisMessage对象
                    RedisMessage redisMessage = JSONUtil.toBean(m, RedisMessage.class);
                    Message message = null;
                    //根据消息类型封装具体消息对象
                    if (redisMessage.getMessageType().equals(MessageType.USER)) {
                        message =  new UserMessage(redisMessage.getMessageType(),redisMessage.getTextContent(),redisMessage.getMedia(),redisMessage.getMetadata());
                    } else if (redisMessage.getMessageType().equals(MessageType.ASSISTANT)) {
                        message =  new AssistantMessage(redisMessage.getTextContent(),redisMessage.getMetadata(),redisMessage.getToolCalls(), redisMessage.getMedia());
                    } else if(redisMessage.getMessageType().equals(MessageType.SYSTEM)) {
                        message = new SystemMessage(redisMessage.getTextContent());
                    } else if(redisMessage.getMessageType().equals(MessageType.TOOL)) {
                        message =  new ToolResponseMessage(redisMessage.getToolResponses(),redisMessage.getMetadata());
                }

                return message;
        }).toList();
        return list;
    }

    /**
     * 清理对话消息
     * @param conversationId
     */
    @Override
    public void clear(String conversationId) {
        // 拼接Rediskey
        String key = getRediskey (conversationId);
        // 删除缓存
        redisTemplate.unlink(key);
    }

    /**
     * 拼接Redis key
     */
    private String getRediskey(String conversationId){
        return "CHAT" + conversationId;
    }
}
